import torch.optim as optim

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 5  # 训练轮数

for epoch in range(num_epochs):
    model.train()  # 设置模型为训练模式
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, labels in train_loader:
        # 将数据移动到GPU（如果可用）
        if torch.cuda.is_available():
            inputs, labels = inputs.cuda(), labels.cuda()

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 统计损失和准确率
        running_loss += loss.item() * inputs.size(0)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = 100. * correct / total

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')


# pip install tensorboard

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/simple_cnn_mnist')  # 指定日志目录

epoch_loss = running_loss / total
epoch_acc = 100. * correct / total

# 记录损失和准确率到TensorBoard
writer.add_scalar('Training/Loss', epoch_loss, epoch)
writer.add_scalar('Training/Accuracy', epoch_acc, epoch)


if torch.cuda.is_available():
    inputs, labels = inputs.cuda(), labels.cuda()


model.eval()  # 设置模型为评估模式
correct = 0
total = 0

with torch.no_grad():
    for inputs, labels in test_loader:
        if torch.cuda.is_available():
            inputs, labels = inputs.cuda(), labels.cuda()
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Test Accuracy: {100. * correct / total:.2f}%')

